#include<iostream>
using namespace std;

/*
extern "C" int func_name(int nparam, void **params, int *ndims, int64_t **shapes, const char **dtypes, void *stream, void *extra);
nparam (int): 输入输出总数。比如算子有2个输入，1个输出，则nparam的值为3。
params (void **): 输入输出指针数组。比如算子有2个输入，1个输出，params[0]指向第一个输入数据，params[1]指向第二个输入数据，params[2]指向输出数据
ndims (int *): 输入输出shape维度数组。比如params[i]是个shape[1024, 1024]的张量，则ndims[i]的值为2。
shapes (int64_t **): 输入输出shape数组。比如params[i]是个shape[1024, 1024]的张量，则shapes[i][0]的值为1024，shapes[i][1]的值为1024。
dtypes (const char **): 输入输出数据类型数组。dtypes里的元素取值可为：
”float32”, “float16”, “float”, “float64”, “int”, “int8”, “int16”, “int32”, “int64”, “uint”, “uint8”, “uint16”, “uint32”, “uint64”, “bool”。
stream (void *): CUDA流指针，仅定义GPU算子实现时需要。
extra (void *): 用于后续扩展。
*/
/*
shapes_backward=((1,1,712,1072),(1072,712),(1,1,712,1072))
types=(ms.float64,ms.int32,ms.float64)
op_backwrad=ops.Custom("./loss_backward.so:CustomBackward", out_shape=shapes_backward, out_dtype=types,func_type="aot")
*/
extern "C" int CustomBackward(int nparam, void **params, int *ndims, 
int64_t **shapes, const char **dtypes, void *stream, void *extra){
    //nparam=4
    //nidms:[4,4,1,4]
    //shapes:shape[0][0-3]=[1,1,712,1072]
    //shapes:shape[1][0-1]=[1,1,712,1072]
    //shapes:shape[2][0]=[2]
    //shapes:shape[3][0-3]=[1,1,712,1072]
    if (nparam!=4) return 1;
    float *prob=static_cast<float *>(params[0]);
    int *indicator=static_cast<int *>(params[1]);
    int *num_area=static_cast<int *>(params[2]);
    float *grads=static_cast<float *>(params[3]);
    
    int length=shapes[0][2]*shapes[0][3];
    //int length=712*1072;
    int num_background_reserved=num_area[1];//保留的背景点数量
    int num_forground=num_area[0];//前景区元素数量
    double beta=num_background_reserved/(num_background_reserved+num_forground);
    /*
    grads=beta*(prob-1),lable[i]==255(indicator[i]==1)
    grads=(1-beta)*prob,lable[i]==255&&II(p[i,j]==1)(indicator[i]==2)
    grads=0,lable[i]==1&&II(p[i,j]==0)(indicator[i]==0)
    */
    for(int i=0;i<length;i++){
        switch (indicator[i])
        {
        case 1:
            //前景区域
            grads[i]=beta*(prob[i]-1);
            break;
        case 2:
            //保留的背景区像素点
            grads[i]=(1-beta)*prob[i];
            break;
        default:
            grads[i]=0;
            break;
        }
    }
    return 0;//仅作为判断函数是否执行正确
}